import cv2
import numpy as np
import matplotlib.pyplot as plt
from shapely.geometry import Polygon
from matplotlib.patches import Polygon as mplPolygon
from mpl_toolkits.mplot3d import Axes3D  # 用于3D可视化
import os

class Visualizer:
    '''
    用于可视化病灶区域（包括2D和3D）
    '''
    
    def __init__(self, base_image_path, validated_lesions, transform_matrices, image_paths_list, mapped_image_polygons, image_resolution=(224, 224)):
        '''
        初始化方法
        输入:
            base_image_path (str): 基底图像的路径
            validated_lesions (list[dict]): 验证后的病灶信息列表（来自 LesionCrossValidator）
            transform_matrices (list[array]): 每张图像的转移矩阵
            image_paths_list (list[str]): 原始图像路径列表（不包括基底图像）
            mapped_image_polygons (list[Polygon]): 每张图像在基底坐标系中的覆盖区域
            image_resolution (tuple): 图像分辨率，例如 (224, 224)
        '''
        self.base_image_path = base_image_path
        self.validated_lesions = validated_lesions
        self.transform_matrices = transform_matrices
        self.image_paths_list = image_paths_list
        self.mapped_image_polygons = mapped_image_polygons  # 包含每张图像的覆盖区域多边形
        self.image_resolution = image_resolution  # 新增的图像分辨率参数

        # 加载并调整基底图像
        self.base_image = cv2.imread(self.base_image_path)
        if self.base_image is None:
            raise ValueError(f"无法加载基底图像：{self.base_image_path}")
        # 调整基底图像的大小
        self.base_image = cv2.resize(self.base_image, self.image_resolution)
        # 获取基底图像的尺寸
        self.base_height, self.base_width = self.base_image.shape[:2]
        
        # 加载并调整其他图像
        self.image_list = []
        for image_path in self.image_paths_list:
            image = cv2.imread(image_path)
            if image is None:
                print(f"无法加载图像：{image_path}")
                self.image_list.append(None)
                continue
            # 调整图像大小
            image = cv2.resize(image, self.image_resolution)
            self.image_list.append(image)
    
    def visualize_2D(self, save_path=None):
        '''
        进行2D可视化，将多个图像拼接在一起，显示验证后的病灶区域
        输入:
            save_path (str, 可选): 如果提供路径，将拼接后的图像保存到指定路径
        输出:
            None
        '''
        # 创建一个空白画布，尺寸足够大以容纳所有变换后的图像
        # 计算画布的尺寸
        all_corners = []
        # 首先处理基底图像的角点
        h, w = self.base_height, self.base_width
        corners = np.array([[0, 0],
                            [w, 0],
                            [w, h],
                            [0, h]], dtype=np.float32).reshape(-1, 1, 2)
        all_corners.extend(corners.reshape(-1, 2))
        
        # 然后处理其他图像的角点
        for idx, image in enumerate(self.image_list):
            if image is None:
                continue
            h, w = image.shape[:2]
            # 获取图像的四个角点
            corners = np.array([[0, 0],
                                [w, 0],
                                [w, h],
                                [0, h]], dtype=np.float32).reshape(-1, 1, 2)
            # 将角点映射到基底坐标系
            M = self.transform_matrices[idx]
            if M is None:
                continue
            transformed_corners = cv2.perspectiveTransform(corners, M)
            all_corners.extend(transformed_corners.reshape(-1, 2))
        
        # 获取所有角点的最小和最大值
        all_corners = np.array(all_corners)
        x_min, y_min = np.int32(all_corners.min(axis=0))
        x_max, y_max = np.int32(all_corners.max(axis=0))
        
        # 计算偏移量
        x_offset = -x_min if x_min < 0 else 0
        y_offset = -y_min if y_min < 0 else 0
        
        # 计算画布尺寸
        canvas_width = x_max - x_min
        canvas_height = y_max - y_min
        
        # 创建画布
        stitched_image = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8)
        
        # 将基底图像放置到画布上
        base_position = (x_offset, y_offset)
        stitched_image[y_offset:y_offset+self.base_height, x_offset:x_offset+self.base_width] = self.base_image

        # 遍历每张图像，进行透视变换并叠加到画布上
        for idx, image in enumerate(self.image_list):
            if image is None:
                continue
            M = self.transform_matrices[idx]
            if M is None:
                continue
            h, w = image.shape[:2]
            # 调整变换矩阵以考虑偏移量
            M_offset = M.copy()
            M_offset[0, 2] += x_offset
            M_offset[1, 2] += y_offset
            # 透视变换图像
            warped_image = cv2.warpPerspective(image, M_offset, (canvas_width, canvas_height))
            # 创建掩码，标记非黑色区域
            mask = (warped_image > 0).any(axis=2)
            # 将图像叠加到画布上
            stitched_image[mask] = warped_image[mask]

        # 绘制验证后的病灶区域
        for lesion in self.validated_lesions:
            lesion_polygon = lesion['polygon']
            # 将多边形坐标转换为画布坐标
            x, y = lesion_polygon.exterior.xy
            x = np.array(x) + x_offset
            y = np.array(y) + y_offset
            pts = np.vstack((x, y)).astype(np.int32).T
            if lesion['status'] == 'True Positive':
                color = (0, 0, 255)  # 红色
            else:
                color = (255, 0, 0)  # 蓝色
            cv2.polylines(stitched_image, [pts], isClosed=True, color=color, thickness=2)

        # 显示或保存结果
        if save_path:
            cv2.imwrite(save_path, stitched_image)
            print(f"拼接后的图像已保存到：{save_path}")
        else:
            plt.figure(figsize=(15, 15))
            plt.imshow(cv2.cvtColor(stitched_image, cv2.COLOR_BGR2RGB))
            plt.axis('off')
            plt.show()

    def visualize_3D(self, save_path=None):
       import cv2
import numpy as np
import matplotlib.pyplot as plt
from shapely.geometry import Polygon
from matplotlib.patches import Polygon as mplPolygon
from mpl_toolkits.mplot3d import Axes3D  # 用于3D可视化
import os

class Visualizer:
    '''
    用于可视化病灶区域（包括2D和3D）
    '''
    
    def __init__(self, base_image_path, validated_lesions, transform_matrices, image_paths_list, mapped_image_polygons, image_resolution=(224, 224)):
        '''
        初始化方法
        输入:
            base_image_path (str): 基底图像的路径
            validated_lesions (list[dict]): 验证后的病灶信息列表（来自 LesionCrossValidator）
            transform_matrices (list[array]): 每张图像的转移矩阵
            image_paths_list (list[str]): 原始图像路径列表（不包括基底图像）
            mapped_image_polygons (list[Polygon]): 每张图像在基底坐标系中的覆盖区域
            image_resolution (tuple): 图像分辨率，例如 (224, 224)
        '''
        self.base_image_path = base_image_path
        self.validated_lesions = validated_lesions
        self.transform_matrices = transform_matrices
        self.image_paths_list = image_paths_list
        self.mapped_image_polygons = mapped_image_polygons  # 包含每张图像的覆盖区域多边形
        self.image_resolution = image_resolution  # 新增的图像分辨率参数

        # 加载并调整基底图像
        self.base_image = cv2.imread(self.base_image_path)
        if self.base_image is None:
            raise ValueError(f"无法加载基底图像：{self.base_image_path}")
        # 调整基底图像的大小
        self.base_image = cv2.resize(self.base_image, self.image_resolution)
        # 获取基底图像的尺寸
        self.base_height, self.base_width = self.base_image.shape[:2]
        
        # 加载并调整其他图像
        self.image_list = []
        for image_path in self.image_paths_list:
            image = cv2.imread(image_path)
            if image is None:
                print(f"无法加载图像：{image_path}")
                self.image_list.append(None)
                continue
            # 调整图像大小
            image = cv2.resize(image, self.image_resolution)
            self.image_list.append(image)
    
    def visualize_2D(self, save_path=None):
        '''
        进行2D可视化，将多个图像拼接在一起，显示验证后的病灶区域
        输入:
            save_path (str, 可选): 如果提供路径，将拼接后的图像保存到指定路径
        输出:
            None
        '''
        # 创建一个空白画布，尺寸足够大以容纳所有变换后的图像
        # 计算画布的尺寸
        all_corners = []
        # 首先处理基底图像的角点
        h, w = self.base_height, self.base_width
        corners = np.array([[0, 0],
                            [w, 0],
                            [w, h],
                            [0, h]], dtype=np.float32).reshape(-1, 1, 2)
        all_corners.extend(corners.reshape(-1, 2))
        
        # 然后处理其他图像的角点
        for idx, image in enumerate(self.image_list):
            if image is None:
                continue
            h, w = image.shape[:2]
            # 获取图像的四个角点
            corners = np.array([[0, 0],
                                [w, 0],
                                [w, h],
                                [0, h]], dtype=np.float32).reshape(-1, 1, 2)
            # 将角点映射到基底坐标系
            M = self.transform_matrices[idx]
            if M is None:
                continue
            transformed_corners = cv2.perspectiveTransform(corners, M)
            all_corners.extend(transformed_corners.reshape(-1, 2))
        
        # 获取所有角点的最小和最大值
        all_corners = np.array(all_corners)
        x_min, y_min = np.int32(all_corners.min(axis=0))
        x_max, y_max = np.int32(all_corners.max(axis=0))
        
        # 计算偏移量
        x_offset = -x_min if x_min < 0 else 0
        y_offset = -y_min if y_min < 0 else 0
        
        # 计算画布尺寸
        canvas_width = x_max - x_min
        canvas_height = y_max - y_min
        
        # 创建画布
        stitched_image = np.zeros((canvas_height, canvas_width, 3), dtype=np.uint8)
        
        # 将基底图像放置到画布上
        base_position = (x_offset, y_offset)
        stitched_image[y_offset:y_offset+self.base_height, x_offset:x_offset+self.base_width] = self.base_image

        # 遍历每张图像，进行透视变换并叠加到画布上
        for idx, image in enumerate(self.image_list):
            if image is None:
                continue
            M = self.transform_matrices[idx]
            if M is None:
                continue
            h, w = image.shape[:2]
            # 调整变换矩阵以考虑偏移量
            M_offset = M.copy()
            M_offset[0, 2] += x_offset
            M_offset[1, 2] += y_offset
            # 透视变换图像
            warped_image = cv2.warpPerspective(image, M_offset, (canvas_width, canvas_height))
            # 创建掩码，标记非黑色区域
            mask = (warped_image > 0).any(axis=2)
            # 将图像叠加到画布上
            stitched_image[mask] = warped_image[mask]

        # 绘制验证后的病灶区域
        for lesion in self.validated_lesions:
            lesion_polygon = lesion['polygon']
            # 将多边形坐标转换为画布坐标
            x, y = lesion_polygon.exterior.xy
            x = np.array(x) + x_offset
            y = np.array(y) + y_offset
            pts = np.vstack((x, y)).astype(np.int32).T
            if lesion['status'] == 'True Positive':
                color = (0, 0, 255)  # 红色
            else:
                color = (255, 0, 0)  # 蓝色
            cv2.polylines(stitched_image, [pts], isClosed=True, color=color, thickness=2)

        # 显示或保存结果
        if save_path:
            cv2.imwrite(save_path, stitched_image)
            print(f"拼接后的图像已保存到：{save_path}")
        else:
            plt.figure(figsize=(15, 15))
            plt.imshow(cv2.cvtColor(stitched_image, cv2.COLOR_BGR2RGB))
            plt.axis('off')
            plt.show()

    def visualize_3D(self, save_path=None):
        '''
        执行3D可视化，将图像和病灶映射到球面上，模拟眼球的形状。
        
        输入:
            save_path (str, 可选): 如果提供路径，将保存静态3D图像到指定路径。
        输出:
            None
        '''
        # 定义代表眼球的球面
        radius = 100  # 设置球面半径
        u = np.linspace(0, 2 * np.pi, 360)
        v = np.linspace(0, np.pi, 180)
        u, v = np.meshgrid(u, v)
        x_sphere = radius * np.sin(v) * np.cos(u)
        y_sphere = radius * np.sin(v) * np.sin(u)
        z_sphere = radius * np.cos(v)

        # 创建3D绘图
        fig = plt.figure(figsize=(14, 14))
        ax = fig.add_subplot(111, projection='3d')

        # 绘制球面
        ax.plot_surface(x_sphere, y_sphere, z_sphere, rstride=10, cstride=10, color='white', alpha=0.1, edgecolor='gray')

        # 收集所有映射后的 x 和 y 坐标，以便归一化
        all_x = []
        all_y = []

        # 遍历每张图像，收集角点坐标
        for idx, (image, M) in enumerate(zip(self.image_list, self.transform_matrices)):
            if M is None:
                continue
            h, w = image.shape[:2]  # 获取图像尺寸
            # 定义图像的四个角点
            image_corners = np.array([[0, 0],
                                      [w, 0],
                                      [w, h],
                                      [0, h]], dtype=np.float32).reshape(-1, 1, 2)
            # 将角点映射到基底坐标系
            mapped_corners = cv2.perspectiveTransform(image_corners, M)
            x_base = mapped_corners[:, 0, 0]
            y_base = mapped_corners[:, 0, 1]
            all_x.extend(x_base)
            all_y.extend(y_base)

        # 收集病灶的坐标
        for lesion in self.validated_lesions:
            lesion_polygon = lesion['polygon']
            x_base, y_base = lesion_polygon.exterior.xy
            all_x.extend(x_base)
            all_y.extend(y_base)

        # 计算归一化所需的最小和最大值
        x_min, x_max = min(all_x), max(all_x)
        y_min, y_max = min(all_y), max(all_y)

        # 再次遍历每张图像，进行归一化并映射到球面
        for idx, (image, M) in enumerate(zip(self.image_list, self.transform_matrices)):
            if M is None:
                continue
            h, w = image.shape[:2]
            image_corners = np.array([[0, 0],
                                      [w, 0],
                                      [w, h],
                                      [0, h]], dtype=np.float32).reshape(-1, 1, 2)
            mapped_corners = cv2.perspectiveTransform(image_corners, M)
            x_base = mapped_corners[:, 0, 0]
            y_base = mapped_corners[:, 0, 1]

            # 归一化坐标
            x_norm = (x_base - x_min) / (x_max - x_min)
            y_norm = (y_base - y_min) / (y_max - y_min)

            # 映射到球面坐标
            theta = x_norm * 2 * np.pi  # 经度
            phi = y_norm * np.pi  # 纬度
            x_mapped = radius * np.sin(phi) * np.cos(theta)
            y_mapped = radius * np.sin(phi) * np.sin(theta)
            z_mapped = radius * np.cos(phi)

            # 绘制图像覆盖区域
            verts = np.vstack((x_mapped, y_mapped, z_mapped)).T
            ax.plot_trisurf(verts[:, 0], verts[:, 1], verts[:, 2], color='white', alpha=0.3)

        # 绘制验证后的病灶区域
        for lesion in self.validated_lesions:
            lesion_polygon = lesion['polygon']
            x_base, y_base = lesion_polygon.exterior.xy

            # 归一化坐标
            x_norm = (np.array(x_base) - x_min) / (x_max - x_min)
            y_norm = (np.array(y_base) - y_min) / (y_max - y_min)

            # 映射到球面坐标
            theta = x_norm * 2 * np.pi
            phi = y_norm * np.pi
            x_mapped = radius * np.sin(phi) * np.cos(theta)
            y_mapped = radius * np.sin(phi) * np.sin(theta)
            z_mapped = radius * np.cos(phi)

            # 绘制病灶点
            ax.scatter(x_mapped, y_mapped, z_mapped, color='red', s=50)

        # 设置绘图参数
        ax.set_xlim(-radius, radius)
        ax.set_ylim(-radius, radius)
        ax.set_zlim(-radius, radius)
        ax.set_box_aspect([1, 1, 1])  # 设置等比例显示
        plt.title('3D视图下的验证病灶区域')

        # 如果提供了保存路径，将图像保存到指定路径
        if save_path:
            plt.savefig(save_path)
        else:
            plt.show()